import unittest

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

'''
Basic version of the prisoner's dilemma.
'''

class PrisonerTestCase(unittest.TestCase):

    def setUp(self):
        #setup world, agent, and basics
        self.maxRounds=1
        self.world = World()

        #add agents        
        self.muscle = Agent('Muscle')
        self.world.addAgent(self.muscle)
        self.brains = Agent('Brains')
        self.world.addAgent(self.brains)

        #order            
        self.world.setOrder([set([self.muscle.name, self.brains.name])])#[self.muscle.name, self.brains.name])
        
    def addStates(self):
        #define states
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)
        
        self.world.defineState(self.muscle.name,'mDefect',bool)
        self.world.setState(self.muscle.name,'mDefect',False)

        self.world.defineState(self.brains.name,'bDefect',bool)
        self.world.setState(self.brains.name,'bDefect',False)

        self.world.defineState(self.muscle.name, 'mReward', int, lo=0, hi=5)
        self.world.setState(self.muscle.name, 'mReward', 3)

        self.world.defineState(self.brains.name, 'bReward', int, lo=0, hi=5)
        self.world.setState(self.brains.name, 'bReward', 3)     
                    
    def addTermination(self):
        #terminate if maxRounds reached
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, False: False}))

    def addActions(self):
        #add actions and appropriate dynamics

        #defect        
        mDefectAction = self.muscle.addAction({'verb': 'defect'})
        tree = makeTree(setTrueMatrix(stateKey(self.muscle.name, 'mDefect')))
        self.world.setDynamics(stateKey(self.muscle.name, 'mDefect'), mDefectAction, tree, enforceMin=True)

        bDefectAction = self.brains.addAction({'verb': 'defect'})
        tree = makeTree(setTrueMatrix(stateKey(self.brains.name, 'bDefect')))
        self.world.setDynamics(stateKey(self.brains.name, 'bDefect'), bDefectAction, tree, enforceMin=True)

        tree = makeTree({'if': equalRow(stateKey(self.muscle.name, 'mDefect'), True),
                         True: setToConstantMatrix(stateKey(self.muscle.name, 'mReward'), 2),
                         False: setToConstantMatrix(stateKey(self.muscle.name, 'mReward'), 1),})
        self.world.setDynamics(stateKey(self.muscle.name,'mReward'), bDefectAction, tree, enforceMin=True)        

        tree = makeTree({'if': equalRow(stateKey(self.muscle.name, 'mDefect'), True),
                         True: setToConstantMatrix(stateKey(self.brains.name, 'bReward'), 2),
                         False: setToConstantMatrix(stateKey(self.brains.name, 'bReward'), 4) })
        self.world.setDynamics(stateKey(self.brains.name,'bReward'), bDefectAction, tree, enforceMin=True)       

        #stay silent
        mSilence = self.muscle.addAction({'verb': 'silent'})
        tree = makeTree(setFalseMatrix(stateKey(self.muscle.name, 'mDefect')))
        self.world.setDynamics(stateKey(self.muscle.name,'mDefect'), mSilence, tree, enforceMin=True)

        bSilence = self.brains.addAction({'verb': 'silent'})
        tree = makeTree(setFalseMatrix(stateKey(self.brains.name, 'bDefect')))
        self.world.setDynamics(stateKey(self.brains.name, 'bDefect'), bSilence, tree, enforceMin=True)

        tree = makeTree({'if': equalRow(stateKey(self.muscle.name, 'mDefect'), True),
                         True: setToConstantMatrix(stateKey(self.muscle.name, 'mReward'), 4),
                         False: setToConstantMatrix(stateKey(self.muscle.name, 'mReward'), 3) })
        
        self.world.setDynamics(stateKey(self.muscle.name,'mReward'), bSilence, tree, enforceMin=True)        

        tree = makeTree({'if': equalRow(stateKey(self.muscle.name, 'mDefect'), True),
                         True: setToConstantMatrix(stateKey(self.brains.name, 'bReward'), 1),
                         False: setToConstantMatrix(stateKey(self.brains.name, 'bReward'), 3) })
        self.world.setDynamics(stateKey(self.brains.name,'bReward'), bSilence, tree, enforceMin=True)
        
    def setModelAndRewards(self):
        #set models, rewards, and beliefs
        self.muscle.setReward(maximizeFeature(stateKey(self.muscle.name,'mReward')), 1.0, True)
        self.brains.setReward(maximizeFeature(stateKey(self.brains.name,'bReward')), 1.0, True)
        
    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()
        self.setModelAndRewards()

    def testRun(self):
        self.testSetup()
        self.world.printState(beliefs=True)
        for t in range(self.maxRounds + 1):
            self.world.explain(self.world.step(),level=2)
            self.world.state.select()
            self.world.printState(beliefs=True)
            if self.world.terminated():
                break

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(PrisonerTestCase):
       if method.startswith("test"):
          suite.addTest(PrisonerTestCase(method))
    unittest.TextTestRunner().run(suite)
            